import { h, makeArray, Schema } from "koishi"

export const name = "dpsk"
export const inject = {
  required: ["database"],
  optional: ["monetary", "autoDeleteResponse"],
}

export const Config = Schema.object({
  cost: Schema.number().default(2),
  reasoningCost: Schema.number().default(6),
})

function argCountAsserter(min, max = min) {
  return ({ options, session, args }) => {
    if (options.help) return
    if (args.length < min) return session.i18n("internal.insufficient-arguments")
    if (args.length > max) return session.i18n("internal.redunant-arguments")
  }
}

export function apply(ctx, config) {
  ctx.model.extend(
    "dpsk.preset",
    {
      name: "string",
      content: "text",
    },
    { primary: "name" },
  )

  ctx.command("dpsk")

  const usersWithPendingRequest = new Set()

  ctx.command("dpsk.ask <prompt:rawtext>", { checkUnknown: true })
    .option("preset", "-p <name>")
    .option("reasoning", "-r")
    .usage(session => session.text("commands.dpsk.ask.cost", config) + "<br/>" + session.text("commands.dpsk.ask.usage"))
    .before(argCountAsserter(1))
    .action(async ({ session, options, root }, prompt) => {
      if (usersWithPendingRequest.has(session.uid)) {
        const r = session.i18n(".pending-request")
        if (session.messageId) r.unshift(h.quote(session.messageId))
        if (ctx.autoDeleteResponse) ctx.autoDeleteResponse.send(session, r)
        else session.send(r)
        return []
      }

      if (ctx.monetary) try {
        await ctx.monetary.cost(session.user.id, options.reasoning ? config.reasoningCost : config.cost)
      } catch (e) {
        if (e?.constructor?.name !== "InsufficientBalanceError") throw e
        return session.i18n("general.insufficient-balance")
      }

      usersWithPendingRequest.add(session.uid)
      try {
        let messages = [{ role: "user", content: prompt }]
        if (options.preset) {
          const [data] = await ctx.database.get("dpsk.preset", options.preset, ["content"])
          if (!data) return session.i18n("commands.dpsk.messages.preset-not-found", [options.preset])
          const { content } = data
          messages.unshift({ role: "system", content })
        }
        ctx.logger.debug(messages)

        const response = await ctx.http.post(
          "https://api.deepseek.com/chat/completions",
          JSON.stringify({
            model: options.reasoning ? "deepseek-reasoner" : "deepseek-chat",
            messages,
            stream: false,
          }),
          {
            headers: {
              "Content-Type": "application/json",
              Accept: "application/json",
              Authorization: "Bearer " + process.env.DEEPSEEK_API_KEY,
            },
          }
        )

        const result = response.choices[0].message
        let outputs = [result.content]
        if (result.reasoning_content) outputs.unshift(result.reasoning_content)

        // add quotes
        if (root && session.messageId)
          outputs = outputs.map(m => [h.quote(session.messageId), m])
        // merge
        if (root && outputs.length === 2 && session.platform === "onebot")
          outputs = [h("message", { forward: true }, outputs[0]), outputs[1]]
        else
          outputs = [outputs[0], ...outputs.slice(1).flatMap(m => [h("message"), m])]

        return outputs.map(h.toElementArray).flat()
      } finally {
        usersWithPendingRequest.delete(session.uid)
      }
    })

  ctx.command("dpsk.preset")

  ctx.command("dpsk.preset.define <name> <content:rawtext>", { strictOptions: true, authority: 2 })
    .option("rest", "-- <content:rawtext>", { hidden: true, descPath: "lnnbot.general.rest-option" })
    .before(argCountAsserter(2))
    .action(async ({ session }, name, content) => {
      if (!name.match(/^[\p{L}\p{N}_-]+$/u)) return session.i18n(".invalid-name")
      await ctx.database.upsert("dpsk.preset", [{ name, content }])
      return session.i18n(".saved", [name])
    })

  ctx.command("dpsk.preset.list", { checkArgCount: true, checkUnknown: true })
    .action(async ({ session }) => {
      const data = await ctx.database.get("dpsk.preset", {}, ["name"])
      if (!data.length) return session.i18n(".none")
      return [
        ...session.i18n(".header"),
        ...data.map(({ name }) => name).reduce((f, i) => makeArray(f).concat(...session.i18n("general.comma"), i)),
      ]
    })

  ctx.command("dpsk.preset.show <name>")
    .before(argCountAsserter(1))
    .action(async ({ session }, name) => {
      const [data] = await ctx.database.get("dpsk.preset", name, ["content"])
      if (!data) return session.i18n("commands.dpsk.messages.preset-not-found", [name])
      return h.text(data.content)
    })

  ctx.command("dpsk.preset.remove <name>", { authority: 2 })
    .before(argCountAsserter(1))
    .action(async ({ session }, name) => {
      await ctx.database.remove("dpsk.preset", name)
      return session.i18n(".removed", [name])
    })

  ctx.i18n.define("zh-CN", "commands.dpsk", {
    description: "DeepSeek 功能",
    messages: {
      "preset-not-found": "预设“{0}”不存在…",
    },
    ask: {
      description: "问 DeepSeek",
      arguments: {
        prompt: "问题",
      },
      cost: "花费：{cost}（不启用深度思考）或 {reasoningCost}（启用深度思考）点数",
      usage: "对话无上下文，每次调用都是新对话。",
      options: {
        preset: "系统级提示词预设名称",
        reasoning: "启用深度思考",
      },
      messages: {
        "pending-request": "您的上一次请求还未完成，请稍后再试…",
      },
    },
    preset: {
      description: "管理系统级提示词预设",
      define: {
        description: "定义预设",
        arguments: {
          name: "要定义的预设名称",
          content: "预设内容",
        },
        messages: {
          "invalid-name": "预设名称无效…",
          saved: "预设“{0}”已保存。",
        },
      },
      list: {
        description: "显示预设列表",
        messages: {
          header: "当前已定义的预设有：",
          none: "当前未定义任何系统级提示词预设。",
        },
      },
      show: {
        description: "查看预设内容",
        arguments: {
          name: "要查看的预设名称",
        },
      },
      remove: {
        description: "移除预设",
        arguments: {
          name: "要移除的预设名称",
        },
        messages: {
          removed: "预设已移除。",
        },
      },
    },
  })
}
